import os
import sys
import glob
import string
import pefile
import hashlib
from operator import itemgetter
from collections import defaultdict, Counter

DEBUG = False
SPACE = "="*20
MALWARE_DIR = "MALWARE/*"  #"JSON/*"
MIN_STRING_LENGTH = 8
common_strings = {}
common_imports = {}
common_md5s = defaultdict(list)
ext = [".sh", ".txt", ".py", ".json", ".pl", ".php"]
logo = """
 __  __    _    _        ____                       _
|  \/  |  / \  | |      / ___|_ __ _   _ _ __   ___| |__   ___ _ __
| |\/| | / _ \ | |     | |   | '__| | | | '_ \ / __| '_ \ / _ \ '__|
| |  | |/ ___ \| |___  | |___| |  | |_| | | | | (__| | | |  __/ |
|_|  |_/_/   \_\_____|  \____|_|   \__,_|_| |_|\___|_| |_|\___|_|
"""


def md5_stats():
    if common_md5s:
        for md5 in common_md5s:
            if len(common_md5s[md5]) > 1:
                print "\nFound A Name/MD5 Collison for '%s'" % md5
                for filename in common_md5s[md5]:
                    print "\t",filename


def string_stats():
    if common_strings:
        print "\n" + SPACE + "50 Most Common Strings" + SPACE
        print "Most Common String", max(common_strings.iteritems(), key=itemgetter(1))
        for i in Counter(common_strings).most_common(50):
            print i


def import_stats():
    if common_imports:
        print "\n" + SPACE + "50 Most Common Imports" + SPACE
        print "Most Common String", max(common_imports.iteritems(), key=itemgetter(1))
        for i in Counter(common_imports).most_common(50):
            print i


def strings(filename):
    x = list(get_strings(filename))
    for i in x:
        if i in common_strings:
            common_strings[i] += 1
        else:
            common_strings[i] = 1
    return x


def get_strings(filename, min=MIN_STRING_LENGTH):
    #http://stackoverflow.com/questions/17195924/python-equivalent-of-unix-strings-utility
    with open(filename, "rb") as f:
        result = ""
        for c in f.read():
            if c in string.printable:
                result += c
                continue
            if len(result) >= min:
                yield result
            result = ""


def process_imports(filename, pestuff):
    if DEBUG:
        print "Imports for %s" % filename
    for entry in pestuff.DIRECTORY_ENTRY_IMPORT:
        if DEBUG:
            print '\t', entry.dll
        for imp in entry.imports:
            if DEBUG:
                print '\t\t', imp.name
            if imp.name is not None:
                if imp.name in common_imports:
                    common_imports[imp.name] += 1
                else:
                    common_imports[imp.name] = 1


def md5_all(dir):
    x = [(f, hashlib.md5(open(f, 'rb').read()).digest()) for f in glob.glob(dir)]
    if DEBUG:
        print x


def md5_one(filename):
    md5 = hashlib.md5(open(filename, 'rb').read()).hexdigest()
    if DEBUG:
        print md5

    common_md5s[md5].append(filename)


def main():
    try:
        files = glob.glob(MALWARE_DIR)
        num_files = len(files)
        proc_files = 0
        print "Number of Files", num_files
        files_left = 0
        for f in files:
            if os.path.splitext(f)[1] in ext:
                continue

            files_left += 1
            sys.stdout.write("\rProcessing [%s of %s][Valid PE: %s]" % (str(files_left), num_files, proc_files))
            sys.stdout.flush()
            if DEBUG:
                print SPACE + "Strings" + SPACE
            strings(f)

            try:
                pe = pefile.PE(f)
                pe.parse_data_directories()
            except Exception, e:
                sys.stderr.write("\nCould not open '%s' for PE processing. Is it a binary?\n" % f)
                continue

            proc_files += 1

            if DEBUG:
                print SPACE + "Imports" + SPACE
            process_imports(f, pe)

            if DEBUG:
                print SPACE + "MD5s" + SPACE
            md5_one(f)

            if DEBUG:
                print SPACE + "Most Common String: " + SPACE
                print max(common_strings.iteritems(), key=itemgetter(1))
    except KeyboardInterrupt:
        string_stats()
        import_stats()
        md5_stats()
        sys.exit(1)
    except Exception, e:
        print "ERROR: ", e
        sys.exit(1)

    string_stats()
    import_stats()
    md5_stats()

    # md5_all(MALWARE_DIR)

if __name__ == "__main__":
    print logo
    main()